
from ast import main
from llama_index.core.indices.struct_store.sql_retriever import SQLRetriever
from llama_index.core.workflow import Workflow, Context, StartEvent, Event, step, StopEvent

from database_rag import DatabaseRAG
from prompts import TEXT_TO_SQL_PROMPT, RESPONSE_SYNTHESIS_PROMPT
from llms import deepseek_llm
from utils import parse_response_to_sql
class TableRetrieveEvent(Event):
    query: str
    table_content_str: str

class TextToSQLEvent(Event):
    query: str
    sql: str

class TextToSQLWorkflow(Workflow, DatabaseRAG):
    def __init__(self, **kwargs):
        Workflow.__init__(self, **kwargs)
        DatabaseRAG.__init__(self, **kwargs)

    @step
    async def retrieve_tables(self, ctx: Context, ev: StartEvent) -> TableRetrieveEvent:
        index = await self.load_index(collection_name="database")
        # 非常关键的步骤
        r = index.as_retriever(similarity_top_k=20)
        table_schema_objs = r.retrieve(ev.query)
        text_list = [node.text for node in table_schema_objs]
        table_content_str = "\n\n".join(text_list)
        return TableRetrieveEvent(query=ev.query, table_content_str=table_content_str)

    @step
    async def generate_sql(self, ctx: Context, ev: TableRetrieveEvent) -> TextToSQLEvent:
        # 替换提示词中的占位符
        format_message = TEXT_TO_SQL_PROMPT.format_messages(
            query=ev.query,
            table_context_str=ev.table_content_str
        )
        llm = deepseek_llm()
        chat_response = llm.chat(format_message)
        sql = parse_response_to_sql(chat_response)
        return TextToSQLEvent(query=ev.query, sql=sql)
    @step
    async def generate_response(self, ctx: Context, ev: TextToSQLEvent) -> StopEvent:
        sql_retriever = SQLRetriever(sql_database=self.sql_database)
        rows = sql_retriever.retrieve(ev.sql)

        # 替换提示词中的占位符
        format_message = RESPONSE_SYNTHESIS_PROMPT.format_messages(
            query=ev.query,
            sql=ev.sql,
            context_str=rows
        )
        # 将加工后的提示词返回前端，前端通过流式调用输出，提升用户体验
        return StopEvent(result={"message": format_message, "sql": ev.sql})

